import os
from moviepy.editor import VideoFileClip
from PIL import Image

# 定义视频目录和输出目录
videos_dir = './videos'
output_dir = './images'

# 创建输出目录（如果不存在）
os.makedirs(output_dir, exist_ok=True)

# 定义截取的时间间隔（单位：秒）
interval = 10

# 遍历videos目录下的所有视频文件
for video_file in os.listdir(videos_dir):
    # 检查文件扩展名是否为视频格式
    if video_file.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
        video_path = os.path.join(videos_dir, video_file)
        print(f"Processing video: {video_file}")
        
        try:
            # 创建视频文件剪辑对象
            video_clip = VideoFileClip(video_path)
            
            # 获取视频的总时长（单位：秒）
            total_duration = int(video_clip.duration)
            
            # 为每个视频创建单独的输出目录
            video_output_dir = os.path.join(output_dir, os.path.splitext(video_file)[0])
            os.makedirs(video_output_dir, exist_ok=True)
            
            # 遍历视频，每隔 interval 秒截取一张图片
            for time_stamp in range(0, total_duration, interval):
                # 获取视频帧
                frame = video_clip.get_frame(time_stamp)
                
                # 将帧转换为 PIL 图像对象
                img = Image.fromarray(frame)
                
                # 保存图像
                output_path = os.path.join(video_output_dir, f"frame_{time_stamp:04d}.png")
                img.save(output_path)
                print(f"Saved frame at {time_stamp}s to {output_path}")
            
            # 释放视频文件
            video_clip.close()
            print(f"Finished processing {video_file}")
            
        except Exception as e:
            print(f"Error processing {video_file}: {str(e)}")
            continue

print("All videos processed successfully!")